# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         utils.py
# Date:           2025-05-24
# Description:    本模块封装了项目中常用的工具函数，支持图像处理、数据保存与匹配算法等功能。
#                 主要功能包括：
#                 - 图像的等比缩放与填充（resize_with_padding）；
#                 - CSV 文件的初始化与数据追加写入（ensure_file, save_to_train_data）；
#                 - 基于匈牙利算法的点集最优匹配（match_points_by_hungarian）；
#                 - 自动收集高质量样本并用于训练数据生成（handle_valid_projection, save_to_train_data）；
#                 是整个骰子姿态识别系统中基础支撑功能的核心工具模块。
# ==============================================================================

import cv2
import numpy as np
import os
import csv

def ensure_file(csv_file):
    # 如果文件不存在，则创建并写入表头
    if not os.path.exists(csv_file):
        with open(csv_file, mode='w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow([
                "img_count",
                "rvec_x", "rvec_y", "rvec_z"
            ])

def handle_valid_projection(mask, rvec, repro_error):
    """
    处理满足最小重投影误差的投影结果，提取轮廓角点并保存到训练数据中。
    """
    from configs.config import Config
    from core.identify_surface import extract_corner_points

    if repro_error < Config.NORMAL_MIN_ALLOW_REPROJECTION_ERROR and Config.SAVE_DATA_TO_TRAIN_FLAG:
        contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        for index, contour in enumerate(contours):
            if hierarchy[0][index][3] == -1:
                corners = extract_corner_points(contour)
                # 保存到训练数据
                save_to_train_data(Config.SAVE_DATA_TO_TRAIN_CSV_FILE, mask, corners, rvec)


# 保存预测结果到CSV
def save_to_train_data(csv_file, section, points, rvec):
    from configs.config import Config

    processed_img = resize_with_padding(section, target_size=128)

    cv2.imwrite(f"reproImages//{Config.IMG_COUNT}.png", processed_img)
    # 处理旋转向量
    rvec_list = rvec.flatten().tolist()
    row = [Config.IMG_COUNT, *rvec_list]

    with open(csv_file, mode='a', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(row)

    print("save_to_train_data", Config.IMG_COUNT, "\t", rvec_list)
    Config.IMG_COUNT += 1


def resize_with_padding(image, target_size=128):
    old_size = image.shape[:2]  # (height, width)
    max_dim = max(old_size)
    scale = target_size / max_dim
    new_size = tuple(int(dim * scale) for dim in old_size[::-1])  # (width, height)

    # 缩放图像
    resized_image = cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)

    # 创建黑色背景的目标图像
    delta_w = target_size - new_size[0]
    delta_h = target_size - new_size[1]
    top, bottom = delta_h // 2, delta_h - delta_h // 2
    left, right = delta_w // 2, delta_w - delta_w // 2

    color = [0, 0, 0] if len(image.shape) == 3 else 0
    resized_image = cv2.copyMakeBorder(resized_image, top, bottom, left, right,
                                       borderType=cv2.BORDER_CONSTANT,
                                       value=color)
    return resized_image

from scipy.optimize import linear_sum_assignment
def match_points_by_hungarian(ref_points, points):
    N = len(ref_points)
    cost_matrix = np.zeros((N, N))

    for i in range(N):
        for j in range(N):
            cost_matrix[i, j] = np.linalg.norm(ref_points[i] - points[j])

    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    return points[col_ind]


def find_best_pose(corners_list, pnp_errors, ba_errors, posenet_errors, rvec_list, tvec_list, ba_rvecs, ba_tvecs,
                   posenet_rvecs):
    """
    查找具有最小重投影误差的最佳姿态估计。
    """
    min_error = float('inf')
    best_idx = -1
    best_method = ""  # 'pnp', 'ba', 'posenet'

    for i in range(len(corners_list)):
        methods = [("pnp", pnp_errors, rvec_list, tvec_list),
                   ("ba", ba_errors, ba_rvecs, ba_tvecs),
                   ("posenet", posenet_errors, posenet_rvecs, tvec_list)]

        for method_name, errors, rvecs, tvecs in methods:
            if errors[i] < min_error:
                min_error = errors[i]
                best_idx = i
                best_method = method_name
                best_rvec = rvecs[best_idx]
                best_tvec = tvecs[best_idx]

    return best_method, best_idx, best_rvec, best_tvec, min_error


def rvec_to_euler(rvec):
    """
    将旋转向量转换为欧拉角。
    """
    R, _ = cv2.Rodrigues(rvec)
    pitch = np.arcsin(R[2, 1])
    yaw = np.arctan2(-R[2, 0], R[2, 2])
    roll = np.arctan2(-R[0, 1], R[1, 1])
    return np.array([pitch, yaw, roll]) * 180 / np.pi